# common/emqxUtil/emqx_client.py
import paho.mqtt.client as mqtt
import time
import json
from typing import Callable, Optional
from datetime import datetime  # 必须添加这行导入

class EMQXClient:
    def __init__(
        self,
        client_id: str,
        broker: str = "192.168.1.42",
        port: int = 8083,
        path: str = "/mqtt",
        username: str = "admin",
        password: str = "yutang1314",
        use_websocket: bool = True
    ):
        """
        EMQX 客户端工具类
        
        :param client_id: 客户端唯一标识
        :param broker: 服务器地址
        :param port: 端口号
        :param path: WebSocket路径
        :param username: 认证用户名
        :param password: 认证密码
        :param use_websocket: 是否使用WebSocket协议
        """
        self.config = {
            "client_id": client_id,
            "broker": broker,
            "port": port,
            "path": path,
            "username": username,
            "password": password,
            "use_websocket": use_websocket
        }
        
        # 回调函数
        self.on_message_callback: Optional[Callable] = None
        self.on_connect_callback: Optional[Callable] = None
        
        # 创建客户端
        transport = "websockets" if use_websocket else "tcp"
        self.client = mqtt.Client(client_id=client_id, transport=transport)
        
        # 初始化配置
        self._setup_client()

    def _setup_client(self):
        """客户端基础配置"""
        if self.config["use_websocket"]:
            self.client.ws_set_options(path=self.config["path"])
            
        self.client.username_pw_set(
            self.config["username"], 
            self.config["password"]
        )
        
        # 绑定默认回调
        self.client.on_connect = self._on_connect
        self.client.on_message = self._on_message
        self.client.on_subscribe = self._on_subscribe
        self.client.on_publish = self._on_publish

    def connect(self, keepalive: int = 60):
        """建立连接"""
        self.client.connect(
            self.config["broker"],
            self.config["port"],
            keepalive
        )
        self.client.loop_start()
        time.sleep(1)  # 等待连接建立

    def disconnect(self):
        """断开连接"""
        self.client.loop_stop()
        self.client.disconnect()

    def publish(self, topic: str, payload, qos: int = 1, retain: bool = False):
        """
        发布消息
        
        :param topic: 主题
        :param payload: 消息内容（支持字符串或字典）
        :param qos: 服务质量等级
        :param retain: 是否保留消息
        """
        if isinstance(payload, dict):
            payload = json.dumps(payload)
            
        result = self.client.publish(topic, payload, qos=qos, retain=retain)
        time.sleep(0.1)  # 确保消息发送
        return result

    def subscribe(self, topic: str, qos: int = 1):
        """订阅主题"""
        self.client.subscribe(topic, qos=qos)

    @staticmethod
    def _get_current_time():
        """获取格式化的当前时间（精确到秒）"""
        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # 回调函数处理 ==============================================
    def _on_connect(self, client, userdata, flags, rc):
        """默认连接回调"""
        if rc == 0:
            print(f"[{self._get_current_time()}] [EMQX] 连接成功 ({self.config['client_id']})")
        else:
            print(f"[{self._get_current_time()}] [EMQX] 连接失败，错误码: {rc}")
            
        if self.on_connect_callback:
            self.on_connect_callback(client, userdata, flags, rc)

    def _on_message(self, client, userdata, msg):
        """默认消息回调"""
        try:
            payload = msg.payload.decode()
            try:
                payload = json.loads(payload)
            except json.JSONDecodeError:
                pass
                
            print(f"[EMQX] 收到消息 ({msg.topic}): {payload}")
        except UnicodeDecodeError:
            print(f"[EMQX] 收到二进制消息 ({msg.topic}): {msg.payload}")
            
        if self.on_message_callback:
            self.on_message_callback(client, userdata, msg)

    @staticmethod
    def _on_subscribe(client, userdata, mid, granted_qos):
        print(f"[EMQX] 订阅确认 (QoS: {granted_qos})")

    @staticmethod
    def _on_publish(client, userdata, mid):
        print(f"[EMQX] 消息发送确认 (ID: {mid})")

    # 回调注册方法 ==============================================
    def register_message_callback(self, callback: Callable):
        """注册自定义消息回调"""
        self.on_message_callback = callback

    def register_connect_callback(self, callback: Callable):
        """注册自定义连接回调"""
        self.on_connect_callback = callback

    # 上下文管理器 ==============================================
    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()